import os
import subprocess
import uuid
from copy import deepcopy

from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric import rsa

from my_models import client, delete_database, save_json, Tickets

config = {
    "IMAGE_NAME": "step3_simulator",  # 镜像名称
    "NUM_DEVICES": 10,  # 启动docker数量
    "NUM_THREADS": 10,  # 每docker模拟数量（线程）
    "SERVER_IP": "127.0.0.1",  # 服务器IP
    "SERVER_PORT": 5000,  # 服务器端口
    "SIMULATE_TIME": 72,  # 模拟时间（小时）
    # 模拟时间加速倍数，1 表示实时模拟
    # 例如 150 表示 1 秒相当于 100 秒模拟时间
    # 假设容器数 * 线程数为100，模拟时间加速倍数为150，后端每秒要处理约25条数据
    # 请根据设备性能调整，建议范围 1-300
    "TIME_SPEEDUP": 300
}

# 获取当前文件路径
base_path = os.path.dirname(__file__)
# 临时文件路径
temp_path = os.path.join(base_path, "temp")


def generate_key():
    """
    生成密钥对
    如果密钥对已存在，则验证密钥对是否正确
    """
    private_key_path = os.path.join(os.path.dirname(base_path), "private_key.pem")
    public_key_path = os.path.join(base_path, "simulator_docker", "public_key.pem")
    # 检查私钥和公钥是否存在
    try:
        # 读取私钥
        f = open(private_key_path, "rb")
        private_key = serialization.load_pem_private_key(f.read(), password=None)
        f.close()
        # 读取公钥
        f = open(public_key_path, "rb")
        public_key = serialization.load_pem_public_key(f.read())
        f.close()
        # 如果读取成功，则验证密钥对
        if private_key.public_key() == public_key:
            return
    except Exception as e:
        # 如果读取或验证失败，则继续重新生成密钥
        print(f"密钥验证失败，错误信息：{e}，重新生成密钥。")
    # 生成 RSA 私钥
    private_key = rsa.generate_private_key(
        public_exponent=65537,
        key_size=2048,
    )
    # 将私钥保存到文件中
    f = open(private_key_path, "wb")
    f.write(
        private_key.private_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PrivateFormat.TraditionalOpenSSL,
            encryption_algorithm=serialization.NoEncryption()
        )
    )
    f.close()
    # 提取公钥,将公钥保存到文件中
    public_key = private_key.public_key()
    f = open(public_key_path, "wb")
    f.write(
        public_key.public_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PublicFormat.SubjectPublicKeyInfo
        )
    )
    f.close()
    print("密钥对已生成并保存")
    input("请重新启动服务器，输入1继续")


def generate_tickets():
    """
    生成 ticket
    """
    # 用于启动docker参数
    device = []
    # 替换宿主机ip
    if "127.0.0.1" in config["SERVER_IP"] or "localhost" in config["SERVER_IP"]:
        config["SERVER_IP"] = "host.docker.internal"
    # 保存 ticket，使用集合保证不重复
    tickets = set()
    for i in range(config["NUM_DEVICES"]):
        device.append(deepcopy(config))
        device[i]["devices"] = []
        for j in range(config["NUM_THREADS"]):
            # device_id
            device_id = "device_" + str(i) + "_" + str(j)
            # 生成 ticket
            ticket = str(uuid.uuid4())
            # 保证不重复
            while ticket in tickets:
                ticket = str(uuid.uuid4())
            tickets.add(ticket)
            # 生成字典
            d = Tickets(device_id=device_id, ticket=ticket).model_dump()
            # 加入启动docker参数
            device[i]["devices"].append(d)
            # 保存到数据库
            client.blood_pressure.tickets.insert_one(deepcopy(d))
    return device


def build_docker():
    """
    检查docker是否构建
    若未构建则提示构建
    """
    # 如果镜像已存在则返回
    result = subprocess.run(['docker', 'images'], capture_output=True, text=True)
    while config["IMAGE_NAME"] not in result.stdout:
        print("cd " + os.path.join(base_path, "simulator_docker"))
        print("sudo docker pull python:3.12-slim")
        print(f"sudo docker build -t {config['IMAGE_NAME']} .")
        input("请在simulator_docker目录下构建镜像，完成后输入1继续")
        result = subprocess.run(['docker', 'images'], capture_output=True, text=True)


def start_dockor(device: list):
    """
    启动docker
    :param device: 启动docker参数
    """
    for i in range(config["NUM_DEVICES"]):
        # 生成配置文件
        file_path = os.path.join(temp_path, f"config_{i}.json")
        # 公钥路径
        public_key_path = os.path.join(base_path, "simulator_docker", "public_key.pem")
        # 模拟程序路径
        simulator_path = os.path.join(base_path, "simulator_docker", "simulator.py")
        save_json(file_path, device[i])
        # docker 名称
        docker_name = f"{config['IMAGE_NAME']}_{i}"
        # 删除已存在的同名容器
        docker_str = f"docker rm -f {docker_name}"
        subprocess.run(docker_str, stderr=subprocess.DEVNULL, stdout=subprocess.DEVNULL)
        # 启动docker
        docker_str = (f"docker run -d "
                      f"--name {docker_name} "  # 容器名称
                      f"-v {file_path}:/app/config.json "  # 配置文件
                      f"-v {public_key_path}:/app/public_key.pem "  # 公钥
                      f"-v {simulator_path}:/app/simulator.py "  # 模拟程序
                      f"{config['IMAGE_NAME']}")  # 镜像名称
        subprocess.run(docker_str)
        print(docker_str)
        print("可通过 docker logs -f " + config["IMAGE_NAME"] + f"_{i} 查看日志")


if __name__ == "__main__":
    delete_database()  # 删除数据库
    generate_key()  # 生成密钥对
    device_main = generate_tickets()  # 生成ticket
    build_docker()  # 检查构建docker
    start_dockor(device_main)  # 启动docker
